#!/usr/bin/python

import unittest
import random

import arithmetic_parser, arithmetic_compiler

class ArithmeticCompilerTests(unittest.TestCase):
  def test_create_empty_function(self):
    func_dict = {'fname': 'func', 'params': [], 'equations': []}
    func = arithmetic_compiler.create_function(func_dict)
    self.assertEquals(func(), {})
    for x in xrange(1, 20):
      params = (0,) * x
      self.assertRaises(TypeError, func, *params)

  def test_create_empty_function_multi_param(self):
    func_dict = {'fname': 'func', 'params': ['a'], 'equations': []}
    func = arithmetic_compiler.create_function(func_dict)
    self.assertEquals(func(1), {'a': 1})
    for x in xrange(20):
      if x == 1:
        continue
      params = (0,) * x
      self.assertRaises(TypeError, func, *params)

    func_dict = {'fname': 'func', 'params': ['a', 'b', 'c'], 'equations': []}
    func = arithmetic_compiler.create_function(func_dict)
    self.assertEquals(func(1, 2, 3), {'a':1, 'b':2, 'c':3})
    for x in xrange(20):
      if x == 3:
        continue
      params = (0,) * x
      self.assertRaises(TypeError, func, *params)

  def test_create_function_not_enough_equations(self):
    func_dict = {'fname': 'func', 'params': [], 'equations':
        [
          ({'a': 1, 'b': 1}, 10),
        ]
    }
    self.assertRaises(arithmetic_compiler.InvalidInputError, arithmetic_compiler.create_function, func_dict)
    func_dict = {'fname': 'func', 'params': ['a'], 'equations':
        [
          ({'a': 1, 'b': 1, 'c': 1}, 10),
        ]
    }
    self.assertRaises(arithmetic_compiler.InvalidInputError, arithmetic_compiler.create_function, func_dict)
    func_dict = {'fname': 'func', 'params': ['a'], 'equations':
        [
          ({'a': 1, 'b': 1, 'c': 1}, 10),
          ({'a': 1, 'd': 2}, 20),
        ]
    }
    self.assertRaises(arithmetic_compiler.InvalidInputError, arithmetic_compiler.create_function, func_dict)

  def test_create_function_contradictions(self):
    func_dict = {'fname': 'func', 'params': [], 'equations':
        [
          ({'a': 1}, 10),
          ({'a': 2}, 10),
        ]
    }
    self.assertRaises(arithmetic_compiler.InvalidInputError, arithmetic_compiler.create_function, func_dict)
    func_dict = {'fname': 'func', 'params': [], 'equations':
        [
          ({'a': 1, 'b': 2}, 10),
          ({'b': 1}, 10),
          ({'a': 1}, 10),
        ]
    }
    self.assertRaises(arithmetic_compiler.InvalidInputError, arithmetic_compiler.create_function, func_dict)
    func_dict = {'fname': 'func', 'params': [], 'equations':
        [
          ({'a': 1, 'b': 2}, 10),
          ({'a': 2, 'b': 4}, 20),
        ]
    }
    self.assertRaises(arithmetic_compiler.InvalidInputError, arithmetic_compiler.create_function, func_dict)
    func_dict = {'fname': 'func', 'params': ['a'], 'equations':
        [
          ({'a': 1}, 10),
        ]
    }
    func = arithmetic_compiler.create_function(func_dict)
    self.assertRaises(arithmetic_compiler.InvalidInputError, func , 1)
    self.assertEquals({'a' : 10}, func(10))
    func_dict = {'fname': 'func', 'params': ['a', 'b', 'c'], 'equations':
        [
          ({'a': 1, 'b': 1, 'c': 1}, 10),
        ]
    }
    func = arithmetic_compiler.create_function(func_dict)
    self.assertRaises(arithmetic_compiler.InvalidInputError, func, 1, 1, 1)
    self.assertEquals({'a' : 3, 'b' : 4, 'c' : 3}, func(3,4,3))
    
  def test_create_function_zero_matrix(self): 
      func_dict = {'fname': 'func', 'params': [], 'equations':
        [
          ({'a': 0, 'b': 0, 'c': 0}, 0),
          ({'a': 0, 'b': 0, 'c': 0}, 0),
          ({'a': 0, 'b': 0, 'c': 0}, 0),
        ]
      }
      self.assertRaises(arithmetic_compiler.InvalidInputError, arithmetic_compiler.create_function, func_dict)
      func_dict = {'fname': 'func', 'params': ['a', 'b', 'c'], 'equations':
        [
          ({'a': 0, 'c': 0, 'e': 0}, 0),
          ({'a': 0, 'b': 0, 'e': 0, 'f': 0}, 0),
          ({'a': 0, 'c': 0, 'd': 0, 'f': 0}, 0),
        ]
      }
      self.assertRaises(arithmetic_compiler.InvalidInputError, arithmetic_compiler.create_function, func_dict)
      
      

  @staticmethod
  def solution_func(a, b, c):
    # This is the solution function for the following matrix:
    #  (1 0 0 0 0 0 | input )
    #  (0 1 0 0 0 0 | input )
    #  (0 0 1 0 0 0 | input )
    #  (1 0 1 0 1 0 | 10    )
    #  (1 1 0 0 1 1 | 20    )
    #  (1 0 1 1 0 1 | 30    )
    d = {'a': a, 'b': b, 'c': c}
    d['e'] = 10 - a - c
    d['f'] = 20 - a - b - d['e']
    d['d'] = 30 - a - c - d['f']
    return d

  @staticmethod
  def float_dict_equals(a, b):
    if set(a.keys()) != set(b.keys()):
      return False
    for key in a.keys():
      if round(a[key], 5) != round(b[key], 5):
        return False
    return True

  def test_create_function(self):
    # This test solves the matrxi from the solution_func()
    func_dict = {'fname': 'func', 'params': ['a', 'b', 'c'], 'equations':
        [
          ({'a': 1, 'c': 1, 'e': 1}, 10),
          ({'a': 1, 'b': 1, 'e': 1, 'f': 1}, 20),
          ({'a': 1, 'c': 1, 'd': 1, 'f': 1}, 30),
        ]
    }
    func = arithmetic_compiler.create_function(func_dict)
    for x in xrange(100):
      r1 = random.random() * 100
      r2 = random.random() * 100
      r3 = random.random() * 100
      expected = ArithmeticCompilerTests.solution_func(r1, r2, r3)
      received = func(r1, r2, r3)
      self.assertTrue(ArithmeticCompilerTests.float_dict_equals(expected, received),
                        "Failure for (%s, %s, %s)\nExpected: %s\nReceived: %s" % (r1, r2, r3, expected, received))


  def test_compile(self):
    # This test solves the matrix from solution_func() for func1, and a very
    # basic one for func2.
    code = """
func1(a, b, c) {
  a + c + e = 10
  a + b - 20 = -e - f
  30 = a + c + d + f
}

func2(input) {
  input + output = 20
}
    """
    compiled = arithmetic_compiler.compile(code)
    self.assertEquals(len(compiled.keys()), 2)
    # Testing func1
    for x in xrange(50):
      r1 = random.random() * 100
      r2 = random.random() * 100
      r3 = random.random() * 100
      expected = ArithmeticCompilerTests.solution_func(r1, r2, r3)
      actual = compiled['func1'](r1, r2, r3)
      self.assertTrue(ArithmeticCompilerTests.float_dict_equals(expected, actual), "Failure for (%s, %s, %s)" % (r1, r2, r3))
    # Testing func2
    for x in xrange(50):
      r = random.random() * 100
      expected =  {'input': r, 'output' : 20 - r}
      actual = compiled['func2'](r)
      self.assertTrue(ArithmeticCompilerTests.float_dict_equals(expected, actual), "Failure for input=" + str(r))


class ArithmeticParserTests(unittest.TestCase):
  def test_parse_empty_func(self):
    code = """empty_func() {}"""
    parsed = arithmetic_parser.parse(code)
    self.assertEquals(len(parsed), 1)
    self.assertEquals(parsed[0]['fname'], 'empty_func')
    self.assertFalse(parsed[0]['params'])
    self.assertFalse(parsed[0]['equations'])

  def test_parse_spacing_between_functions(self):
    code = """
first_func() {} SECOND_FUNCTION() {}






third0123456789Function() {}
    """
    parsed = arithmetic_parser.parse(code)
    self.assertEquals(len(parsed), 3)
    for x in xrange(3):
      self.assertFalse(parsed[x]['params'])
      self.assertFalse(parsed[x]['equations'])
    self.assertEquals(parsed[0]['fname'], 'first_func')
    self.assertEquals(parsed[1]['fname'], 'SECOND_FUNCTION')
    self.assertEquals(parsed[2]['fname'], 'third0123456789Function')

  def test_parse_different_paramlists(self):
    code = """
empty_func() {}
empty_func_newline_empyt_paramlist(
) {}
empty_func_one_param(a) {}
empty_func_mult_param(firstParam, second_param, Third0Param) {}
empty_func_strange_param_spaces(   a   ,   b

,

c


) {}
    """
    parsed = arithmetic_parser.parse(code)
    self.assertEquals(len(parsed), 5)
    for x in xrange(5):
      self.assertFalse(parsed[x]['equations'])

    self.assertEquals(parsed[0]['fname'], 'empty_func')
    self.assertFalse(parsed[0]['params'])

    self.assertEquals(parsed[1]['fname'], 'empty_func_newline_empyt_paramlist')
    self.assertFalse(len(parsed[1]['params']))

    self.assertEquals(parsed[2]['fname'], 'empty_func_one_param')
    self.assertEquals(len(parsed[2]['params']), 1)
    self.assertEquals(parsed[2]['params'][0], 'a')

    self.assertEquals(parsed[3]['fname'], 'empty_func_mult_param')
    self.assertEquals(len(parsed[3]['params']), 3)
    self.assertEquals(parsed[3]['params'][0], 'firstParam')
    self.assertEquals(parsed[3]['params'][1], 'second_param')
    self.assertEquals(parsed[3]['params'][2], 'Third0Param')

    self.assertEquals(parsed[4]['fname'], 'empty_func_strange_param_spaces')
    self.assertEquals(len(parsed[4]['params']), 3)
    self.assertEquals(parsed[4]['params'][0], 'a')
    self.assertEquals(parsed[4]['params'][1], 'b')
    self.assertEquals(parsed[4]['params'][2], 'c')

  def test_parse_equations_count(self):
    code = """
firstFunc() { a = 1 }
secondFunc() { a = 1
b = 2


c = 3
d = 4
e = 5 }
    """
    parsed = arithmetic_parser.parse(code)
    self.assertEquals(len(parsed), 2)
    for x in xrange(2):
      self.assertFalse(parsed[x]['params'])

    self.assertEquals(parsed[0]['fname'], 'firstFunc')
    self.assertEquals(len(parsed[0]['equations']), 1)

    self.assertEquals(parsed[1]['fname'], 'secondFunc')
    self.assertEquals(len(parsed[1]['equations']), 5)

  def test_parse_equations(self):
    code = """
func1() {
  a=10
  a    =    10
  a    =10
  a - 10 = 0
  a-10=0
  a -10 = 0
  0 = 10 - a
  0 = -a + 10
  a = 5 + 5
  a - 5 = 5
  a = 2*5
  2*a = 20
  2*(a-10) = 0
  2 * (a - 10) = 0
  2 * a - 20 = 0
}

func2() {
  2*(a + b - 12) - c + 15 = -3*(d + 17) - e + 20
}
    """
    parsed = arithmetic_parser.parse(code)
    self.assertEquals(len(parsed), 2)
    for x in xrange(2):
      self.assertFalse(parsed[x]['params'])

    self.assertEquals(parsed[0]['fname'], 'func1')
    #print parsed[0]
    self.assertEquals(len(parsed[0]['equations']), 15)
    equasion_possibilities = (({'a': 1}, 10), ({'a': -1}, -10), ({'a': 2}, 20))
    for x in xrange(15):
      self.assertTrue(parsed[0]['equations'][x] in equasion_possibilities)

    self.assertEquals(parsed[1]['fname'], 'func2')
    self.assertEquals(len(parsed[1]['equations']), 1)
    self.assertEquals(parsed[1]['equations'][0], ({'a': 2, 'b' : 2, 'c': -1, 'd': 3, 'e' : 1}, -22))


if __name__ == '__main__':
  import unittest
  unittest.main()
